Dice roll simulation

Time: O(MxN); Space: O(M); medium

A die simulator generates a random number from 1 to 6 for each roll. You introduced a constraint to the generator such that it cannot roll the number i more than rollMax[i] (1-indexed) consecutive times.

Given an array of integers rollMax and an integer n, return the number of distinct sequences that can be obtained with exact n rolls.

Two sequences are considered different if at least one element differs from each other. Since the answer may be too large, return it modulo 10^9 + 7.

Example 1:

Input: n = 2, rollMax = [1,1,2,2,2,3]

Output: 34

Explanation:

  • There will be 2 rolls of die

  • if there are no constraints on the die, there are 6 * 6 = 36 possible combinations.

  • In this case, looking at rollMax array, the numbers 1 and 2 appear at most once consecutively,

  • therefore sequences (1,1) and (2,2) cannot occur, so the final answer is 36-2 = 34.

Example 2:

Input: n = 2, rollMax = [1,1,1,1,1,1]

Output: 30

Example 3:

Input: n = 3, rollMax = [1,1,1,2,2,3]

Output: 181

Constraints:

  • 1 <= n <= 5000

  • len(rollMax) == 6

  • 1 <= rollMax[i] <= 15

Hints:

  1. Think on Dynamic Programming.

  2. DP(pos, last) which means we are at the position pos having as last the last character seen.

1. Dynamic programming [O(MxN), O(M)]

[3]:
import functools

class Solution1(object):
    """
    Time: O(M*N), M is the max of rollMax
    Space: O(M)
    """
    def dieSimulator(self, n, rollMax):
        """
        :type n: int
        :type rollMax: List[int]
        :rtype: int
        """
        MOD = 10**9+7
        def sum_mod(array):
            return functools.reduce(lambda x, y: (x+y)%MOD, array)

        dp = [[1] + [0]*(rollMax[i]-1) for i in range(6)]      # 0-indexed

        for _ in range(n-1):
            new_dp = [[0]*rollMax[i] for i in range(6)]
            for i in range(6):
                for k in range(rollMax[i]):
                    for j in range(6):
                        if i == j:
                            if k < rollMax[i]-1:  # 0-indexed
                                new_dp[j][k+1] = (new_dp[j][k+1]+dp[i][k])%MOD
                        else:
                            new_dp[j][0] = (new_dp[j][0]+dp[i][k])%MOD
            dp = new_dp

        return sum_mod(sum_mod(row) for row in dp)
[4]:
s = Solution1()

n = 2
rollMax = [1,1,2,2,2,3]
assert s.dieSimulator(n, rollMax) == 34

n = 2
rollMax = [1,1,1,1,1,1]
assert s.dieSimulator(n, rollMax) == 30

n = 3
rollMax = [1,1,1,2,2,3]
assert s.dieSimulator(n, rollMax) == 181